﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;

namespace BciUtils
{
    public class EegDataReader
    {
        private const float restGap = 0.1f;

        private readonly string filename;
        private readonly int epochTime;
        private readonly int restTime;
        private readonly int samplingRate;
        private readonly string[] classes;

        private StreamReader infoReader;
        private StreamReader labelsReader;
        private string dataFileName;

        private int currentPosition;

        private string[] dataLines;

        public EegDataReader(string filename, int epochTime, int restTime, int samplingRate, string[] classes)
        {
            this.filename = filename;
            this.epochTime = epochTime;
            this.restTime = restTime;
            this.samplingRate = samplingRate;
            this.classes = classes;

            infoReader = new StreamReader(filename + EegFileConstants.InfoPostfix);
            labelsReader = new StreamReader(filename + EegFileConstants.LabelsPostfix);
            dataFileName = filename + EegFileConstants.DataPostfix;
        }

        public List<EegLabeledData> GetEegLabeledData(int windowSize, float gap, bool includeRestSamples = false)
        {
            int restGapSize = (int) (restGap * restTime * samplingRate);

            List<EegLabeledData> data = new List<EegLabeledData>();

            int gapSize = (int) (epochTime * samplingRate * gap);

            dataLines = File.ReadAllLines(dataFileName);
                        
            string classLabelString;
            float[] previousValues = null;
            while ((classLabelString = labelsReader.ReadLine()) != null)
            {
                string[] parts = classLabelString.Split(EegFileConstants.ClassLabelSeparator);

                if (parts.Length != 2)
                {
                    throw new ArgumentException("Wrong class label format");
                }

                int classLabel = Array.IndexOf(classes, parts[0]);
                int labelNumber = Int32.Parse(parts[1]);

                int firstSampleNumber = labelNumber + gapSize;
                int lastSampleNumber = firstSampleNumber + epochTime * samplingRate - gapSize;

                int currentSampleNumber = firstSampleNumber;
                List<float[]> windowData = new List<float[]>();
                while (currentSampleNumber + windowSize <= lastSampleNumber)
                {
                    float[] values = ReadValues(currentSampleNumber);

                    if (previousValues != null && previousValues.Length != values.Length)
                    {
                        throw new ArgumentException("Incorrect number of values");
                    }
                    previousValues = values;
      
                    windowData.Add(values);

                    if (windowData.Count == windowSize)
                    {
                        EegLabeledData eegData = new EegLabeledData
                        {
                            Data = windowData,
                            ClassLabel = classLabel,
                            SampleNumber = currentSampleNumber
                            
                        };
                        data.Add(eegData);
                        windowData = new List<float[]>();
                    }

                    currentSampleNumber++;
                }

                if (includeRestSamples)
                {
                    int firtsRestSampleNumber = labelNumber + epochTime * samplingRate + restGapSize;
                    int lastRestSampleNumber = firtsRestSampleNumber + restTime * samplingRate - restGapSize;
                    int currentRestSampleNumber = firtsRestSampleNumber;
                    List<float[]> windowRestData = new List<float[]>();
                    while (currentRestSampleNumber + windowSize <= lastRestSampleNumber)
                    {
                        float[] values = ReadValues(currentRestSampleNumber);

                        windowRestData.Add(values);

                        if (windowRestData.Count == windowSize)
                        {
                            EegLabeledData eegData = new EegLabeledData
                            {
                                Data = windowRestData,
                                ClassLabel = classes.Length
                            };
                            data.Add(eegData);
                            windowRestData = new List<float[]>();
                        }

                        currentRestSampleNumber++;
                    } 
                }
            } 
          
            return data;
        }

        private float[] ReadValues(int sampleNumber)
        {
            string line = dataLines[sampleNumber - 1];
            string[] stringValues = line.Split(EegFileConstants.ValuesSeparator);
            float[] values = stringValues.Select(o => Single.Parse(o)).ToArray();

            Console.WriteLine(sampleNumber + " sample was read...");

            return values;
        }
    }
}
